#
# This file is a part of the open-eBackup project.
# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
# If a copy of the MPL was not distributed with this file, You can obtain one at
# http://mozilla.org/MPL/2.0/.
#
# Copyright (c) [2024] Huawei Technologies Co.,Ltd.
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
#

import os
import time

from common.common_models import LogDetail
from common.const import SubJobStatusEnum, DBLogLevel, ReportDBLabel, AuthType, RestoreTypeEnum
from sqlserver import log
from sqlserver.commons.common import output_action_result, get_key_value, get_restore_channel, is_new_restore_path
from sqlserver.commons.const import SQLServerStrConstant
from sqlserver.commons.const import SubJobExecResult, SQLServerProgressFileType, Progress
from sqlserver.commons.const import VDIActionType, ParamConstant, BodyErr, ProgressConst, \
    SqlServerNameConst, SqlSeverErrorKeyword, SqlServerReportLabel, SQLServerCode
from sqlserver.commons.sqlserver_utils import SqlServerUtils, stop_mssql_instance, start_mssql_instance, \
    get_format_restore_time
from sqlserver.sqlserver_restore_base import SQLServerRestoreBase


class SQLServerRestoreInstance(SQLServerRestoreBase):
    def __init__(self, p_id, job_id, sub_job_id, json_param):
        super().__init__(p_id, job_id, sub_job_id, json_param)
        self._job_id = job_id
        self.sub_job_id = sub_job_id
        self.m_json_param = json_param
        self.m_instance_name = ''
        self.m_default_name = "MSSQLSERVER"
        self.real_restore_databases = []
        self.database_num = 0
        self.error_logtail = LogDetail(logInfo=SqlServerReportLabel.SUB_JOB_FAILED_LABEL,
                                       logDetail=BodyErr.ERROR_INTERNAL.value,
                                       logInfoParam=[self.sub_job_id], logLevel=DBLogLevel.ERROR.value)
        self.success_logtail = LogDetail(logInfo=SqlServerReportLabel.RESTORE_SUB_JOB_SUCCESS_LABEL,
                                         logInfoParam=[self.sub_job_id], logLevel=DBLogLevel.INFO.value)
        self.auth_mode = ""
        self.user_name = ""
        self.pass_prefix = ""
        self.is_new_restore = False

    def get_restore_list(self, data_path):
        if len(self.real_restore_databases) != 0:
            self.database_num = len(self.real_restore_databases)
            return True
        if not os.path.exists(data_path):
            return False
        report_list = SqlServerUtils.get_report_service_database_info(self.m_instance_name)
        for _, _, files in os.walk(data_path):
            for file_mbr in files:
                file_name, file_extension = os.path.splitext(os.path.basename(file_mbr))
                log.info(f'get restore file_name: {file_name}, file_extension: {file_extension}')
                if self.is_new_restore and file_name in SQLServerStrConstant.SYS_DATABASE:
                    continue
                if file_name in self.real_restore_databases:
                    continue
                if file_name in report_list:
                    continue
                self.real_restore_databases.append(file_name)
            break
        self.database_num = len(self.real_restore_databases)
        return True

    def get_restore_database_list(self):
        """
        获取待恢复的数据库列表，如果为空，表示实例全部恢复
        :return:
        """
        restore_sub_object = self.m_json_param.get("job", {}).get("restoreSubObjects", [])
        for mbr in restore_sub_object:
            if self.is_new_restore and mbr.get('name').strip('/') in SQLServerStrConstant.SYS_DATABASE:
                continue
            self.real_restore_databases.append(mbr.get('name').strip('/'))

    def init_param(self):
        flags, self.auth_mode = get_key_value(f'job_targetObject_auth_authType_{self._p_id}')
        if not flags or self.auth_mode == "0":
            flags, self.auth_mode = get_key_value(f'job_targetEnv_nodes_0_auth_authType_'
                                                  f'{self._p_id}')
            if not flags or self.auth_mode == "0":
                return False
        self.m_instance_name = self.m_json_param.get("job", {}).get("targetObject", {}).get("name")
        if str(self.auth_mode) == str(AuthType.APP_PASSWORD.value):
            flags, self.user_name = get_key_value(f'job_targetObject_auth_authKey_{self._p_id}')
            self.pass_prefix = f'job_targetObject_auth_authPwd_{self._p_id}'
            if not flags:
                flags, self.user_name = get_key_value(f'job_targetEnv_nodes_0_auth_authKey_{self._p_id}')
                self.pass_prefix = f'job_targetEnv_nodes_0_auth_authPwd_{self._p_id}'
                if not flags:
                    return False
        self.check_and_prepare_user_info(self.m_instance_name)
        self.new_restore_path = self.sql_utils.get_restore_data_path()
        self.get_restore_new_path()
        return True

    def allow_restore_in_local_node(self):
        log.info("Enter allow_restore_in_local_node_full...")
        self.init_param()
        # 检查实例运行状态
        if not SqlServerUtils.check_service_is_running(self.m_instance_name):
            node_ip = self._json_param.get("job", {}).get("targetEnv", {}).get("nodes", [{}])[0].get("endpoint", "")
            output_action_result(self._p_id, SQLServerCode.FAILED.value,
                                 BodyErr.INSTANCE_STATE_ABNORMAL.value, "", [node_ip, self.m_instance_name])
            return False
        flags = self.sql_utils.check_user_permission(self.auth_mode, self.user_name)
        if not flags:
            output_action_result(self._p_id, SQLServerCode.FAILED.value,
                                 BodyErr.SQLSERVER_PERMISSIONS_ERROR.value, "")
            return False
        # 检查副本
        copy_info = self.get_extend_info()
        src_version = copy_info.get("version", "")
        if not self.sql_utils.check_version_is_matched(src_version):
            output_action_result(self._p_id, SQLServerCode.FAILED.value,
                                 BodyErr.RESTORE_FAIL_FOR_VERSION_INCONSISTENCY.value, "")
            log.error(f"Version is different and can not restore, src_version:{src_version}")
            return False
        log.info("Leave allow_restore_in_local_node_full...")
        return True

    def check_is_same_database_name(self, instance_database_list):
        """
        当细粒度恢复时，检查是否有同名数据库
        """
        is_new_restore = is_new_restore_path(self._json_param)
        if not is_new_restore:
            return True
        for dt_name in instance_database_list:
            if dt_name in self.real_restore_databases and dt_name not in SQLServerStrConstant.SYS_DATABASE:
                log.error(f"Database: {dt_name} is already exist")
                return False
        return True

    def check_is_same_name_in_ag(self, ag_database_list):
        """
        判断待恢复的数据库是否和可用性组的数据库重名
        :return:
        """
        for database in ag_database_list:
            if database in self.real_restore_databases:
                log.error(f"Database: {database} is already in availability_groups")
                return False
        return True

    def check_file_exists(self, file_path, file_list):
        # 只检查待恢复数据库
        meta_info = self.get_meta_info(file_path)
        check_info = []
        for database in self.real_restore_databases:
            check_info.extend(meta_info.get(database, []))
        for info in check_info:
            if info.get("physical_name") in file_list:
                self.restore_fail_tail.log_detail = BodyErr.RESTORE_FAIL_FOR_DATABASE_FILE_EXIST.value
                log.error(f"File: {info.get('physical_name')} exist")
                raise Exception("File exist")

    def check_same_files(self, new_path):
        if not os.path.exists(new_path):
            return True
        target_file_list = os.listdir(new_path)
        # 检查全量副本文件是否存在
        self.check_file_exists(self.full_copy.get("metaPath"), target_file_list)
        # 检查差异副本文件是否存在
        if self.diff_copy:
            self.check_file_exists(self.diff_copy.get("metaPath"), target_file_list)
        # 检查日志副本文件是否存在
        for copy in self.log_copy.get("dataPath"):
            self.check_file_exists(copy, target_file_list)
        return True

    def restore_pre(self):
        log.info("Enter sqlserver pre....")
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, ProgressConst.MIN_PROGRESS.value,
                                    "", SQLServerProgressFileType.COMMON)
        flags = self.init_param()
        if not flags:
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                        "", SQLServerProgressFileType.COMMON)
            log.error("Init param failed")
            return False
        is_new_restore = is_new_restore_path(self._json_param)
        self.is_new_restore = is_new_restore
        self.get_restore_database_list()
        self.get_copy_info()
        flags = self.get_restore_list(self.full_copy.get("dataPath"))
        if not flags:
            log.error(f'Get real restore list failed')
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                        "", SQLServerProgressFileType.COMMON)
            return False
        instance_databases = self.sql_utils.get_instance_databases()
        ag_databases = self.sql_utils.get_ag_database_list()
        # 检查数据库是否同名，是否已存在可用性组
        if not self.check_is_same_database_name(instance_databases) or not self.check_is_same_name_in_ag(ag_databases):
            self.error_logtail.log_info = ReportDBLabel.PRE_REQUISIT_FAILED
            self.error_logtail.log_detail = BodyErr.RESTORE_FAIL_FOR_DATABASE_EXIST.value
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                        self.error_logtail, SQLServerProgressFileType.COMMON)
            return False
        # 数据库级新位置恢复，检查数据库文件是否已存在
        log.info(f"is_new_restore:{is_new_restore}")
        if is_new_restore:
            self.check_same_files(self.new_restore_path)
        self.write_progress_to_file(SubJobStatusEnum.COMPLETED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                    "", SQLServerProgressFileType.COMMON)
        log.info("Leave sqlserver pre....")
        return True

    def exec_restore_pre_job(self):
        try:
            return self.restore_pre()
        except Exception as exception_str:
            log.error(f"restore pre error,{exception_str}")
            self.restore_fail_tail.log_info = ReportDBLabel.PRE_REQUISIT_FAILED
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED,
                                        self.restore_fail_tail, SQLServerProgressFileType.COMMON)
            return False

    def report_running_progress(self, progress):
        SubJobExecResult.SUCCESS_NUM += 1
        self.running_tail.log_info_param = [
            self._sub_job_id, f"{self.database_num - SubJobExecResult.SUCCESS_NUM}", f"{SubJobExecResult.SUCCESS_NUM}"
        ]
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, progress,
                                    self.running_tail, SQLServerProgressFileType.COMMON)

    def master_restore(self, copy_path):
        """
        全量恢复master数据库
        master只支持全量恢复
        """
        log.info("Start stop sqlserver agent")
        ret = stop_mssql_instance(self.m_instance_name)
        if not ret:
            return False
        log.info("Start switch sqlserver agent to single user")
        ret = start_mssql_instance(self.m_instance_name, self.auth_mode, self.user_name, True)
        if not ret:
            return False
        flag = self.do_restore("master", copy_path, copy_path, VDIActionType.FULL_RESTORE_WITH_REPLACE.value, "")
        if not flag:
            return False
        log.info("Succeed to restore master")
        self.report_running_progress(Progress.START_PROGRESS)
        return True

    def restore_master(self, copy_path):
        try:
            ret = self.master_restore(copy_path)
        finally:
            log.info("Start switch sqlserver agent to multi user")
            stop_mssql_instance(self.m_instance_name)
            log.info("Restart sqlserver agent")
            start_res = start_mssql_instance(self.m_instance_name, self.auth_mode, self.user_name, False)
            if not start_res:
                ret = False
        log.info(f"Restore master ret: {ret}")
        return ret

    def do_restore(self, database_name, data_path, meta_path, restore_type, restore_time):
        log.info(f'Start restore database: {database_name}, restore_type: {restore_type}')
        if database_name == SqlServerNameConst.TEMPDB:
            return True
        time.sleep(3)
        is_new_restore = is_new_restore_path(self._json_param)
        if is_new_restore:
            restore_info = self.get_database_restore_info(database_name, self.new_restore_path, meta_path,
                                                          database_name)
        else:
            restore_info = ""
        all_copy_path = os.path.join(data_path, f'{database_name}.bak')
        channel = get_restore_channel(data_path, database_name)
        restore_cmd = f'{ParamConstant.VDI_TOOL_PATH} {restore_type} \"{database_name}\"' \
                      f' \"{all_copy_path}\" \"{self.vdi_info}\" \"{restore_time}\" \"{restore_info}\" {channel}'
        log.info(f"restore_cmd:{restore_cmd}")
        if database_name in ("msdb", "model"):
            for _ in range(3):
                ret, std_out, std_err = self.sql_utils.get_command_result(restore_cmd)
                log.info(f"ret:{ret},std_out:{std_out},std_err:{std_err}")
                if SqlSeverErrorKeyword.INUSE in std_out:
                    raise Exception(f'{database_name} is in use')
                if int(ret) == 0:
                    log.info(f'Succeed to restore database: {database_name}, restore_type: {restore_type}')
                    return True
                time.sleep(60)
            log.info(f'error to restore database: {database_name}, restore_type: {restore_type}')
            return False
        else:
            ret, std_out, std_err = self.sql_utils.get_command_result(restore_cmd)
            log.info(f"ret:{ret},std_out:{std_out},std_err:{std_err}")
            if SqlSeverErrorKeyword.INUSE in std_out:
                raise Exception(f'{database_name} is in use')
            if int(ret) != 0:
                log.info(f'error to restore database: {database_name}, restore_type: {restore_type}')
                return False
            log.info(f'Succeed to restore database: {database_name}, restore_type: {restore_type}')
            return True

    def full_restore_main_process(self):
        process_num = Progress.START_PROGRESS
        for databases_name in self.real_restore_databases:
            if databases_name == SqlServerNameConst.MASTER:
                continue
            log.info(f"Start to restore database: {databases_name}")
            process_num += Progress.INTERVAL_PROGRESS
            if process_num > Progress.MAX_PROGRESS:
                process_num = Progress.MAX_PROGRESS
            if not self.sql_utils.set_database_offline(databases_name):
                return False
            flags = self.do_restore(databases_name, self.full_copy.get("dataPath"), self.full_copy.get("metaPath"),
                                    VDIActionType.FULL_RESTORE_WITH_REPLACE.value, "")
            if not flags:
                return False
            log.info(f"Succeed to restore database: {databases_name}")
            self.report_running_progress(process_num)
        self.success_logtail.log_info_param = [self._sub_job_id, f"{self.database_num}"]
        self.write_progress_to_file(SubJobStatusEnum.COMPLETED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                    self.success_logtail, SQLServerProgressFileType.COMMON)
        return True

    def full_restore(self):
        log.info('Enter sqlserver full restore....')
        self.running_tail.log_info_param = [self._sub_job_id, f"{self.database_num}", "0"]
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, Progress.START_PROGRESS,
                                    self.running_tail, SQLServerProgressFileType.COMMON)
        # 首先恢复master 数据库
        if "master" in self.real_restore_databases:
            flags = self.restore_master(self.full_copy.get("dataPath"))
            if not flags:
                return False
        return self.full_restore_main_process()

    def diff_restore_main_process(self):
        process_num = Progress.START_PROGRESS
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, process_num,
                                    self.success_logtail, SQLServerProgressFileType.COMMON)
        for databases_name in self.real_restore_databases:
            if databases_name == SqlServerNameConst.MASTER:
                continue
            log.info(f"Start to restore database: {databases_name}")
            process_num += Progress.INTERVAL_PROGRESS
            if process_num > Progress.MAX_PROGRESS:
                process_num = Progress.MAX_PROGRESS
            if not self.sql_utils.set_database_offline(databases_name):
                return False
            # 全量恢复
            flags = self.do_restore(databases_name, self.full_copy.get("dataPath"), self.full_copy.get("metaPath"),
                                    VDIActionType.DIFF_RESTORE_WITH_NO_RECOVERY_REPLACE, "")
            if not flags:
                return False
            # 差异恢复
            flags = self.do_restore(databases_name, self.diff_copy.get("dataPath"), self.diff_copy.get("metaPath"),
                                    VDIActionType.DIFF_RESTORE_WITH_RECOVERY, "")
            if not flags:
                return False
            log.info(f"Succeed to restore database: {databases_name}")
            self.report_running_progress(process_num)
        self.success_logtail.log_info_param = [self._sub_job_id, f"{self.database_num}"]
        self.write_progress_to_file(SubJobStatusEnum.COMPLETED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                    self.success_logtail, SQLServerProgressFileType.COMMON)
        log.info('Leave sqlserver diff restore....')
        return True

    def diff_restore(self):
        log.info('Enter sqlserver diff restore....')
        self.running_tail.log_info_param = [self._sub_job_id, f"{self.database_num}", "0"]
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, Progress.START_PROGRESS,
                                    self.running_tail, SQLServerProgressFileType.COMMON)
        # 首先恢复master数据库
        if "master" in self.real_restore_databases:
            flags = self.restore_master(self.diff_copy.get("dataPath"))
            if not flags:
                return False
        return self.diff_restore_main_process()

    def log_restore_other_database(self):
        process_num = Progress.START_PROGRESS
        last_copy_num = len(self.log_copy.get("dataPath")) - 1
        restore_time = self._json_param.get("job", {}).get("extendInfo", {}).get("restoreTimestamp", "")
        if restore_time:
            copies = self._json_param.get("job", {}).get("copies", [])
            restore_time = get_format_restore_time(restore_time, copies)
        for databases_name in self.real_restore_databases:
            if databases_name == SqlServerNameConst.MASTER:
                continue
            log.info(f"Start to restore database: {databases_name}")
            process_num += Progress.INTERVAL_PROGRESS
            if process_num > Progress.MAX_PROGRESS:
                process_num = Progress.MAX_PROGRESS
            if not self.sql_utils.set_database_offline(databases_name):
                return False
            # 全量恢复，命令类型值为7
            flags = self.do_restore(databases_name, self.full_copy.get("dataPath"), self.full_copy.get("metaPath"),
                                    VDIActionType.DIFF_RESTORE_WITH_NO_RECOVERY_REPLACE, "")
            if not flags:
                return False
            # 再做差异恢复, PID恢复可以没有差异副本
            if self.diff_copy:
                flags = self.do_restore(databases_name, self.diff_copy.get("dataPath"), self.diff_copy.get("metaPath"),
                                        VDIActionType.LOG_RESTORE_WITH_NO_RECOVERY, "")
                if not flags:
                    return False
            # 再根据日志备份做类型为9 和10 的日志恢复
            for k, log_path in enumerate(self.log_copy.get("dataPath")):
                flags = self.do_restore(databases_name, log_path, log_path,
                                        VDIActionType.POINT_IN_TIME_RESTORE_WITH_NO_RECOVERY, restore_time)
                if k == last_copy_num:
                    flags = self.do_restore(databases_name, log_path, log_path,
                                            VDIActionType.POINT_IN_TIME_RESTORE_WITH_RECOVERY, restore_time)
                if not flags:
                    return False
            log.info(f"Succeed to restore database: {databases_name}")
            self.report_running_progress(process_num)
        self.success_logtail.log_info_param = [self._sub_job_id, f"{self.database_num}"]
        self.write_progress_to_file(SubJobStatusEnum.COMPLETED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                    self.success_logtail, SQLServerProgressFileType.COMMON)
        return True

    def log_restore(self):
        # PIT恢复
        log.info('Enter sqlserver log restore....')
        self.running_tail.log_info_param = [self._sub_job_id, f"{self.database_num}", "0"]
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, Progress.START_PROGRESS,
                                    self.running_tail, SQLServerProgressFileType.COMMON)
        # 首先恢复master 数据库
        if SqlServerNameConst.MASTER in self.real_restore_databases:
            mst_bak_path = self.full_copy.get("dataPath")
            if self.diff_copy.get("dataPath"):
                mst_bak_path = self.diff_copy.get("dataPath")
            flags = self.restore_master(mst_bak_path)
            if not flags:
                return False
        # 再次恢复其他数据库
        flags = self.log_restore_other_database()
        if not flags:
            return False
        return True

    def restore(self):
        log.info("Enter sqlserver restore....")
        is_new_restore = is_new_restore_path(self._json_param)
        self.is_new_restore = is_new_restore
        self.get_restore_database_list()
        self.get_copy_info()
        self.write_progress_to_file(SubJobStatusEnum.RUNNING.value, ProgressConst.MIN_PROGRESS, "",
                                    SQLServerProgressFileType.COMMON)
        self.init_param()
        flags = self.get_restore_list(self.full_copy.get("dataPath"))
        if not flags:
            log.error("Failed to get restore database")
            return False
        try:
            Progress.INTERVAL_PROGRESS = int(Progress.MAX_PROGRESS / self.database_num)
        except ZeroDivisionError as err:
            raise Exception("The divisor is 0.") from err
        if self.restore_type == RestoreTypeEnum.FULL_RESTORE.value:
            flags = self.full_restore()
        elif self.restore_type == RestoreTypeEnum.DIFF_RESTORE.value:
            flags = self.diff_restore()
        elif self.restore_type == RestoreTypeEnum.LOG_RESTORE.value:
            flags = self.log_restore()
        else:
            log.error(f'Unsupported restore type: {self.restore_type}')
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                        self.error_logtail, SQLServerProgressFileType.COMMON)
            return False
        if not flags:
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED.value,
                                        self.error_logtail, SQLServerProgressFileType.COMMON)
            return False
        log.info("Leave sqlserver restore....")
        return True

    def exec_sub_job(self):
        log_detail = LogDetail(logInfo=ReportDBLabel.RESTORE_SUB_FAILED, logInfoParam=[self._sub_job_id],
                               logLevel=DBLogLevel.ERROR.value)
        try:
            self.restore()
        except Exception as e_info:
            log.exception("Exception when restore")
            str_err = str(e_info)
            if SqlSeverErrorKeyword.INUSE in str_err:
                log_detail.log_detail = BodyErr.SQLSERVER_DATABASE_IN_USE_ERROR.value
                space_index = str_err.find(' ')
                dt_name = str_err[:space_index]
                log_detail.log_detail_param = [dt_name]
            else:
                log_detail.log_detail = BodyErr.ERROR_INTERNAL.value
            self.write_progress_to_file(SubJobStatusEnum.FAILED.value, ProgressConst.PROGRESS_ONE_HUNDRED, log_detail,
                                        SQLServerProgressFileType.COMMON)
            return

    def exec_restore_post(self):
        log.debug("Enter sqlserver restore....")
        self.write_progress_to_file(SubJobStatusEnum.COMPLETED.value, ProgressConst.PROGRESS_ONE_HUNDRED, "",
                                    SQLServerProgressFileType.COMMON)
        log.debug("Leave sqlserver restore....")
